"""
多个SAC数据时间同步模块
用于将多个SAC文件的时间对齐，截取共同时间部分
"""

import os
import numpy as np
from obspy import UTCDateTime, read
from obspy.core import Stream, Trace


def get_sac_time_range(sac_file):
    """
    获取SAC文件的时间范围
    
    参数:
        sac_file (str): SAC文件路径
        
    返回:
        tuple: (开始时间, 结束时间)
    """
    st = read(sac_file)
    if len(st) == 0:
        raise ValueError(f"无法读取SAC文件: {sac_file}")
    
    tr = st[0]
    start_time = tr.stats.starttime
    end_time = tr.stats.endtime
    
    return start_time, end_time


def find_common_time_range(sac_files):
    """
    找到多个SAC文件的共同时间范围
    
    参数:
        sac_files (list): SAC文件路径列表
        
    返回:
        tuple: (共同开始时间, 共同结束时间)
    """
    if not sac_files:
        raise ValueError("SAC文件列表为空")
    
    # 获取所有文件的时间范围和采样率
    time_ranges = []
    sampling_rates = []
    
    for sac_file in sac_files:
        start_time, end_time = get_sac_time_range(sac_file)
        st = read(sac_file)
        if len(st) == 0:
            continue
        sampling_rate = 1.0 / st[0].stats.delta
        
        time_ranges.append((start_time, end_time))
        sampling_rates.append(sampling_rate)
        
        print(f"文件 {sac_file}: 开始时间={start_time}, 结束时间={end_time}, 采样率={sampling_rate}Hz")
    
    # 找到最晚的开始时间和最早的结束时间
    common_start = max([tr[0] for tr in time_ranges])
    common_end = min([tr[1] for tr in time_ranges])
    
    # 检查是否有共同时间范围
    if common_start >= common_end:
        raise ValueError("SAC文件没有共同的时间范围")
    
    print(f"共同时间范围: {common_start} - {common_end}")
    print(f"各文件采样率: {sampling_rates}")
    
    return common_start, common_end


def synchronize_sac_files(sac_files, output_dir, prefix="sync"):
    """
    同步多个SAC文件，截取共同时间部分
    
    参数:
        sac_files (list): SAC文件路径列表
        output_dir (str): 输出目录
        prefix (str): 输出文件名前缀
        
    返回:
        list: 同步后的SAC文件路径列表
    """
    # 创建输出目录
    os.makedirs(output_dir, exist_ok=True)
    
    # 找到共同时间范围
    common_start, common_end = find_common_time_range(sac_files)
    
    # 计算理论采样点数
    duration = common_end - common_start
    
    # 读取所有文件，找到最低采样率
    if not sac_files:
        raise ValueError("SAC文件列表为空")
    
    # 获取所有文件的采样率，找到最低采样率
    sampling_rates = []
    for sac_file in sac_files:
        st = read(sac_file)
        if len(st) == 0:
            continue
        sampling_rate = 1.0 / st[0].stats.delta
        sampling_rates.append(sampling_rate)
    
    if not sampling_rates:
        raise ValueError("无法读取任何SAC文件的采样率")
    
    # 使用最低采样率作为参考
    min_sampling_rate = min(sampling_rates)
    reference_dt = 1.0 / min_sampling_rate
    reference_npts = int(round(duration / reference_dt)) + 1
    
    print(f"使用最低采样率: {min_sampling_rate}Hz")
    print(f"理论采样点数计算: 时长={duration}s, 采样间隔={reference_dt}s, 采样点数={reference_npts}")
    
    # 读取所有文件并存储在内存中
    stream_list = []
    for sac_file in sac_files:
        st = read(sac_file)
        if len(st) == 0:
            print(f"警告: 无法读取SAC文件 {sac_file}，跳过")
            continue
        stream_list.append((sac_file, st))
    
    # 同步每个SAC文件
    synced_files = []
    for sac_file, st in stream_list:
        tr = st[0]
        
        # 确保采样率一致（重采样到最低采样率）
        if tr.stats.delta != reference_dt:
            print(f"重采样 {sac_file} 从 {1.0/tr.stats.delta}Hz 到 {min_sampling_rate}Hz")
            tr.interpolate(sampling_rate=min_sampling_rate)
        
        # 计算精确的结束时间，确保采样点数一致
        exact_end_time = common_start + (reference_npts - 1) * reference_dt
        
        # 截取共同时间部分，使用精确的采样点数控制
        tr.trim(starttime=common_start, endtime=exact_end_time, nearest_sample=True, pad=False)
        
        # 验证采样点数是否一致
        if tr.stats.npts != reference_npts:
            print(f"警告: {sac_file} 采样点数不一致: 期望={reference_npts}, 实际={tr.stats.npts}")
            
            # 更专业的处理方式：使用线性插值
            print(f"  使用线性插值调整到目标采样点数")
            
            # 创建新的时间序列
            original_time = np.arange(0, tr.stats.npts) * tr.stats.delta
            target_time = np.arange(0, reference_npts) * tr.stats.delta
            
            # 使用线性插值
            from scipy import interpolate
            if len(tr.data) > 1:
                interp_func = interpolate.interp1d(original_time, tr.data, 
                                                  kind='linear', 
                                                  fill_value='extrapolate')
                tr.data = interp_func(target_time)
            else:
                # 如果数据太少，使用最后一个值填充
                last_value = tr.data[-1] if len(tr.data) > 0 else 0
                tr.data = np.full(reference_npts, last_value)
            
            tr.stats.npts = reference_npts
        
        # 生成输出文件名
        base_name = os.path.splitext(os.path.basename(sac_file))[0]
        output_file = os.path.join(output_dir, f"{prefix}_{base_name}.sac")
        
        # 保存同步后的SAC文件
        tr.write(output_file, format='SAC')
        synced_files.append(output_file)
        
        print(f"已同步 {sac_file} -> {output_file}")
        print(f"  采样点数: {tr.stats.npts}, 采样率: {1.0/tr.stats.delta}Hz")
    
    return synced_files


def batch_synchronize_sac(
    sac_files: list,
    output_dir: str,
    prefix: str = "sync",
    target_sampling_rate: float | None = None
) -> list:
    """批量同步目录中的SAC文件。

    参数:
        sac_files (list): SAC文件路径列表
        output_dir (str): 输出目录
        prefix (str): 输出文件名前缀
        target_sampling_rate (float | None): 同步后的目标采样率 (Hz)。
            • 若为 None 或 <= 0，则自动采用所有文件中的最低采样率。
            • 若指定，则所有文件将重采样到此采样率。

    返回:
        list: 同步后的SAC文件路径列表
    """
    if not sac_files:
        raise ValueError("SAC文件列表为空")

    print(f"找到 {len(sac_files)} 个SAC文件进行同步")

    # 按分量分组 (X, Y, Z)
    component_groups = {}
    for sac_file in sac_files:
        base_name = os.path.splitext(os.path.basename(sac_file))[0]
        if base_name.endswith('_X') or base_name.endswith('_Y') or base_name.endswith('_Z'):
            common_part = base_name[:-2]
            component = base_name[-1]
            if common_part not in component_groups:
                component_groups[common_part] = {}
            component_groups[common_part][component] = sac_file

    # 找到所有完整分量组的共同时间范围
    all_complete_group_files = []
    for common_part, components in component_groups.items():
        if all(c in components for c in ['X', 'Y', 'Z']):
            all_complete_group_files.extend([components['X'], components['Y'], components['Z']])

    if not all_complete_group_files:
        print("警告: 没有找到任何一个完整的X、Y、Z分量组，将尝试同步所有找到的文件。")
        all_sac_files_to_process = sac_files
    else:
        all_sac_files_to_process = all_complete_group_files

    if not all_sac_files_to_process:
        raise ValueError("没有可供同步的SAC文件")

    # 找到所有文件的共同时间范围
    common_start, common_end = find_common_time_range(all_sac_files_to_process)

    # 计算理论采样点数
    duration = common_end - common_start

    # 确定参考采样率
    if target_sampling_rate is not None and target_sampling_rate > 0:
        reference_sampling_rate = target_sampling_rate
        print(f"使用用户指定的目标采样率: {reference_sampling_rate}Hz")
    else:
        # 获取所有文件的采样率，找到最低采样率
        sampling_rates = []
        for sac_file in all_sac_files_to_process:
            st = read(sac_file)
            if len(st) == 0:
                continue
            sampling_rate = 1.0 / st[0].stats.delta
            sampling_rates.append(sampling_rate)

        if not sampling_rates:
            raise ValueError("无法读取任何SAC文件的采样率")

        reference_sampling_rate = min(sampling_rates)
        print(f"使用最低采样率: {reference_sampling_rate}Hz")

    reference_dt = 1.0 / reference_sampling_rate
    reference_npts = int(round(duration / reference_dt)) + 1

    print(f"理论采样点数计算: 时长={duration}s, 采样间隔={reference_dt}s, 采样点数={reference_npts}")

    # 同步所有找到的文件
    os.makedirs(output_dir, exist_ok=True)

    synced_files = synchronize_sac_files_with_reference(
        all_sac_files_to_process,
        output_dir,
        prefix,
        common_start,
        common_end,
        reference_npts,
        reference_dt
    )

    return synced_files


def synchronize_sac_files_with_reference(sac_files, output_dir, prefix, common_start, common_end, reference_npts, reference_dt):
    """
    使用参考参数同步多个SAC文件
    
    参数:
        sac_files (list): SAC文件路径列表
        output_dir (str): 输出目录
        prefix (str): 输出文件名前缀
        common_start: 共同开始时间
        common_end: 共同结束时间
        reference_npts: 参考采样点数
        reference_dt: 参考采样间隔
        
    返回:
        list: 同步后的SAC文件路径列表
    """
    # 创建输出目录
    os.makedirs(output_dir, exist_ok=True)
    
    # 读取所有文件并存储在内存中
    stream_list = []
    for sac_file in sac_files:
        st = read(sac_file)
        if len(st) == 0:
            print(f"警告: 无法读取SAC文件 {sac_file}，跳过")
            continue
        stream_list.append((sac_file, st))
    
    # 同步每个SAC文件
    synced_files = []
    for sac_file, st in stream_list:
        tr = st[0]
        
        # 备份原始头信息
        original_header = tr.stats.copy()

        # 确保采样率一致
        if abs(tr.stats.delta - reference_dt) > 1e-9: # 避免浮点数精度问题
            print(f"重采样 {sac_file} 从 {tr.stats.sampling_rate:.2f}Hz 到 {1.0/reference_dt:.2f}Hz")
            tr.interpolate(sampling_rate=1.0/reference_dt, starttime=tr.stats.starttime)

        # 计算精确的结束时间，确保采样点数一致
        exact_end_time = common_start + (reference_npts - 1) * reference_dt

        # 截取共同时间部分
        tr.trim(starttime=common_start, endtime=exact_end_time, nearest_sample=True, pad=True, fill_value=0)

        # 验证并修正采样点数
        if tr.stats.npts != reference_npts:
            print(f"警告: {sac_file} 采样点数不一致: 期望={reference_npts}, 实际={tr.stats.npts}。将进行修正。")
            # 创建一个新的Trace对象以确保数据和头信息匹配
            new_data = np.zeros(reference_npts)
            if len(tr.data) > 0:
                # 填充数据，避免越界
                copy_len = min(len(tr.data), reference_npts)
                new_data[:copy_len] = tr.data[:copy_len]
            
            tr.data = new_data
            tr.stats.npts = reference_npts

        # 将原始头信息与更新后的关键信息合并
        final_header = original_header
        final_header.npts = tr.stats.npts
        final_header.delta = tr.stats.delta
        final_header.starttime = tr.stats.starttime
        # endtime, b, e 等字段将由 obspy 根据以上信息自动计算，无需手动设置

        # 创建一个新的Trace对象来保存
        synced_trace = Trace(data=tr.data, header=final_header)

        # 生成输出文件名
        base_name = os.path.splitext(os.path.basename(sac_file))[0]
        output_file = os.path.join(output_dir, f"{prefix}_{base_name}.sac")

        # 保存同步后的SAC文件
        synced_trace.write(output_file, format='SAC')
        synced_files.append(output_file)
        
        print(f"已同步 {sac_file} -> {output_file}")
        print(f"  采样点数: {tr.stats.npts}, 采样率: {1.0/tr.stats.delta}Hz")
    
    return synced_files


if __name__ == "__main__":
    import argparse
    
    parser = argparse.ArgumentParser(description="同步多个SAC文件的时间")
    parser.add_argument("input", help="输入SAC文件或目录")
    parser.add_argument("output", help="输出目录")
    parser.add_argument("--pattern", default="*.sac", help="文件名模式（仅当输入为目录时使用）")
    parser.add_argument("--prefix", default="sync", help="输出文件名前缀")
    
    args = parser.parse_args()
    
    if os.path.isfile(args.input):
        # 输入是单个文件，需要指定其他文件
        print("错误: 当输入为单个文件时，请使用其他方式指定所有需要同步的SAC文件")
    elif os.path.isdir(args.input):
        # 输入是目录，批量处理
        try:
            synced_files = batch_synchronize_sac(args.input, args.output, args.pattern, args.prefix)
            print(f"\n成功同步 {len(synced_files)} 个SAC文件")
        except Exception as e:
            print(f"错误: {e}")
    else:
        print(f"错误: 输入路径 {args.input} 不存在")